# app/api/upload/routes.py

import os
import uuid
import time
from typing import List
from fastapi import APIRouter, File, UploadFile, HTTPException, Form

from config.settings import settings
from app.api.upload.schemas import (
    UploadResponse, 
    TaskStatusResponse, 
    BatchUploadRequest,
    UploadStats
)
from app.api.upload.utils import validate_file_extension, get_file_type, save_upload_file
from app.core.oss_service import oss_service

# 配置日志
from app.core.logging import get_logger
logger = get_logger(__name__)

router = APIRouter()

# 存储任务状态的字典（在生产环境中应该使用Redis或数据库）
task_status = {}

def process_file_async(file_path: str, file_type: str, knowledge_base: str = "default", task_id: str = None, oss_url: str = None):
    """异步处理单个文件"""
    logger.info(f"开始异步处理文件: {file_path} (类型: {file_type}, 知识库: {knowledge_base})")
    start_time = time.time()
    
    try:
        # 更新任务状态为处理中
        if task_id:
            task_status[task_id] = {
                'status': settings.TASK_STATUS_PROCESSING,
                'progress': 0,
                'message': '开始处理文件...',
                'file_path': file_path,
                'file_type': file_type,
                'knowledge_base': knowledge_base
            }
        
        # 导入必要的模块
        from app.core.database import connect_milvus, init_collections
        from app.processors.text_processor import TextProcessor
        from app.processors.image_processor import ImageProcessor
        from app.processors.video_processor import VideoProcessor
        from app.processors.document_processor import DocumentProcessor
        
        # 更新进度
        if task_id:
            task_status[task_id]['progress'] = 10
            task_status[task_id]['message'] = '连接数据库...'
        
        logger.info("连接Milvus数据库...")
        connect_milvus()
        
        # 更新进度
        if task_id:
            task_status[task_id]['progress'] = 20
            task_status[task_id]['message'] = '初始化集合...'
        
        logger.info("初始化集合...")
        init_collections()
        
        # 根据文件类型选择处理器
        logger.info(f"选择处理器: {file_type}")
        if task_id:
            task_status[task_id]['progress'] = 30
            task_status[task_id]['message'] = f'选择处理器: {file_type}'
        
        if file_type in ['pdf', 'docx', 'pptx']:
            processor = DocumentProcessor()
        elif file_type in ['jpg', 'jpeg', 'png', 'bmp', 'gif']:
            processor = ImageProcessor()
        elif file_type in ['mp4', 'avi', 'mov', 'mkv']:
            processor = VideoProcessor()
        else:
            processor = TextProcessor()
        
        # 更新进度
        if task_id:
            task_status[task_id]['progress'] = 50
            task_status[task_id]['message'] = '开始处理文件...'
        
        logger.info("开始处理文件...")
        result = processor.process(file_path, file_type, knowledge_base, oss_url)
        
        processing_time = time.time() - start_time
        logger.info(f"文件处理完成: {file_path} - 耗时: {processing_time:.2f}s")
        
        # 处理完成后删除本地文件
        try:
            if os.path.exists(file_path):
                os.remove(file_path)
                logger.info(f"删除本地文件: {file_path}")
        except Exception as e:
            logger.error(f"删除本地文件 {file_path} 时出错: {e}")
        
        # 更新任务状态为完成
        if task_id:
            task_status[task_id] = {
                'status': settings.TASK_STATUS_COMPLETED,
                'progress': 100,
                'message': '文件处理完成',
                'file_path': file_path,
                'file_type': file_type,
                'knowledge_base': knowledge_base,
                'text_chunks_count': result.get('text_chunks_count', 0),
                'images_count': result.get('images_count', 0),
                'processing_time': processing_time,
                'result': result
            }
        
        return {
            'status': settings.TASK_STATUS_COMPLETED,
            'file_path': file_path,
            'file_type': file_type,
            'text_chunks_count': result.get('text_chunks_count', 0),
            'images_count': result.get('images_count', 0),
            'processing_time': processing_time,
            'message': '文件处理完成'
        }
        
    except Exception as e:
        logger.error(f"处理文件 {file_path} 时发生错误: {str(e)}")
        import traceback
        logger.error(traceback.format_exc())
        
        # 处理失败后也删除本地文件
        try:
            if os.path.exists(file_path):
                os.remove(file_path)
                logger.info(f"删除本地文件: {file_path}")
        except Exception as del_e:
            logger.error(f"删除本地文件 {file_path} 时出错: {del_e}")
        
        # 更新任务状态为失败
        if task_id:
            task_status[task_id] = {
                'status': settings.TASK_STATUS_FAILED,
                'progress': 0,
                'message': f'文件处理失败: {str(e)}',
                'file_path': file_path,
                'file_type': file_type,
                'knowledge_base': knowledge_base,
                'error': str(e),
                'processing_time': time.time() - start_time
            }
        
        return {
            'status': settings.TASK_STATUS_FAILED,
            'file_path': file_path,
            'file_type': file_type,
            'error': str(e),
            'message': f'文件处理失败: {str(e)}',
            'processing_time': time.time() - start_time
        }

@router.post("/single", response_model=UploadResponse, summary="单文件上传")
async def upload_single_file(
    file: UploadFile = File(...),
    knowledge_base: str = Form("default")
):
    """上传单个文件并异步处理"""
    logger.info(f"收到文件上传请求: {file.filename}, 知识库: {knowledge_base}")
    
    try:
        # 验证文件类型
        logger.info(f"验证文件类型: {file.filename}")
        if not validate_file_extension(file.filename):
            logger.error(f"不支持的文件类型: {file.filename}")
            raise HTTPException(
                status_code=400, 
                detail=f"不支持的文件类型: {file.filename}"
            )
        
        # 验证文件大小
        logger.info(f"验证文件大小: {file.size} bytes")
        if file.size and file.size > settings.MAX_FILE_SIZE:
            logger.error(f"文件大小超过限制: {file.size} > {settings.MAX_FILE_SIZE}")
            raise HTTPException(
                status_code=400,
                detail=f"文件大小超过限制: {file.size} > {settings.MAX_FILE_SIZE}"
            )
        
        # 保存文件到本地和阿里云OSS
        logger.info("保存上传文件到本地和阿里云OSS...")
        save_result = save_upload_file(file, knowledge_base)
        local_path = save_result['local_path']
        file_type = save_result['file_type']
        oss_result = save_result['oss_result']
        
        logger.info(f"文件保存成功: {local_path}")
        if oss_result['success']:
            logger.info(f"文件上传到OSS成功: {oss_result['url']}")
        else:
            logger.warning(f"文件上传到OSS失败: {oss_result['error']}")
        
        # 获取文件扩展名
        file_ext = os.path.splitext(file.filename)[1].lower()[1:]  # 去掉点号
        logger.info(f"文件类型: {file_type}, 扩展名: {file_ext}")
        
        # 生成任务ID
        task_id = str(uuid.uuid4())
        
        # 初始化任务状态
        task_status[task_id] = {
            'status': settings.TASK_STATUS_PENDING,
            'progress': 0,
            'message': '文件上传成功，等待处理...',
            'file_path': local_path,
            'file_type': file_ext,
            'knowledge_base': knowledge_base,
            'oss_url': oss_result.get('url', ''),
            'oss_success': oss_result['success']
        }
        
        # 异步启动处理任务
        import threading
        def async_process():
            process_file_async(local_path, file_ext, knowledge_base, task_id, oss_result.get('url'))
        
        thread = threading.Thread(target=async_process)
        thread.daemon = True
        thread.start()
        
        logger.info(f"文件上传成功，任务ID: {task_id}")
        
        # 立即返回成功响应
        return UploadResponse(
            task_id=task_id,
            message=f"文件 {file.filename} 上传成功，正在后台处理中。",
            status="uploaded",
            knowledge_base=knowledge_base
        )
        
    except Exception as e:
        logger.error(f"文件上传处理异常: {str(e)}")
        raise HTTPException(status_code=500, detail=str(e))

@router.post("/batch", response_model=UploadResponse, summary="批量文件上传")
async def upload_batch_files(
    files: List[UploadFile] = File(...),
    knowledge_base: str = Form("default")
):
    """批量上传文件并异步处理"""
    logger.info(f"收到批量文件上传请求: {len(files)} 个文件")
    
    try:
        if not files:
            logger.error("没有选择文件")
            raise HTTPException(status_code=400, detail="没有选择文件")
        
        file_paths = []
        uploaded_files = []
        task_ids = []
        
        for i, file in enumerate(files):
            logger.info(f"处理第 {i+1}/{len(files)} 个文件: {file.filename}")
            
            # 验证文件类型
            if not validate_file_extension(file.filename):
                logger.error(f"不支持的文件类型: {file.filename}")
                raise HTTPException(
                    status_code=400,
                    detail=f"不支持的文件类型: {file.filename}"
                )
            
            # 验证文件大小
            if file.size and file.size > settings.MAX_FILE_SIZE:
                logger.error(f"文件 {file.filename} 大小超过限制: {file.size} > {settings.MAX_FILE_SIZE}")
                raise HTTPException(
                    status_code=400,
                    detail=f"文件 {file.filename} 大小超过限制: {file.size} > {settings.MAX_FILE_SIZE}"
                )
            
            # 保存文件到本地和阿里云OSS
            save_result = save_upload_file(file, knowledge_base)
            local_path = save_result['local_path']
            file_type = save_result['file_type']
            oss_result = save_result['oss_result']
            
            file_paths.append(local_path)
            uploaded_files.append(file.filename)
            
            # 生成任务ID
            task_id = str(uuid.uuid4())
            task_ids.append(task_id)
            
            # 获取文件扩展名
            file_ext = os.path.splitext(file.filename)[1].lower()[1:]  # 去掉点号
            
            # 初始化任务状态
            task_status[task_id] = {
                'status': settings.TASK_STATUS_PENDING,
                'progress': 0,
                'message': f'文件 {file.filename} 上传成功，等待处理...',
                'file_path': local_path,
                'file_type': file_ext,
                'knowledge_base': knowledge_base,
                'oss_url': oss_result.get('url', ''),
                'oss_success': oss_result['success']
            }
            
            # 异步启动处理任务
            import threading
            def async_process():
                process_file_async(local_path, file_ext, knowledge_base, task_id, oss_result.get('url'))
            
            thread = threading.Thread(target=async_process)
            thread.daemon = True
            thread.start()
        
        logger.info(f"批量文件上传成功，生成了 {len(task_ids)} 个任务")
        
        return UploadResponse(
            task_id=",".join(task_ids),  # 多个任务ID用逗号分隔
            message=f"成功上传 {len(files)} 个文件，正在后台处理中。",
            status="uploaded",
            knowledge_base=knowledge_base
        )
        
    except Exception as e:
        logger.error(f"批量文件上传处理异常: {str(e)}")
        raise HTTPException(status_code=500, detail=str(e))

@router.post("/paths", response_model=UploadResponse, summary="通过文件路径批量上传")
async def upload_by_paths(request: BatchUploadRequest):
    """通过文件路径批量上传（适用于服务器本地文件）"""
    logger.info(f"收到路径上传请求: {len(request.file_paths)} 个文件")
    
    try:
        if not request.file_paths:
            logger.error("文件路径列表为空")
            raise HTTPException(status_code=400, detail="文件路径列表为空")
        
        task_ids = []
        
        # 验证文件是否存在
        for i, file_path in enumerate(request.file_paths):
            logger.info(f"验证第 {i+1}/{len(request.file_paths)} 个文件: {file_path}")
            
            if not os.path.exists(file_path):
                logger.error(f"文件不存在: {file_path}")
                raise HTTPException(
                    status_code=400,
                    detail=f"文件不存在: {file_path}"
                )
            
            if not validate_file_extension(file_path):
                logger.error(f"不支持的文件类型: {file_path}")
                raise HTTPException(
                    status_code=400,
                    detail=f"不支持的文件类型: {file_path}"
                )
            
            # 生成任务ID
            task_id = str(uuid.uuid4())
            task_ids.append(task_id)
            
            # 获取文件类型
            file_type = get_file_type(file_path)
            file_ext = os.path.splitext(file_path)[1].lower()[1:]  # 去掉点号
            
            # 上传文件到阿里云OSS
            oss_result = oss_service.upload_file(file_path, file_type, request.knowledge_base)
            
            # 初始化任务状态
            task_status[task_id] = {
                'status': settings.TASK_STATUS_PENDING,
                'progress': 0,
                'message': f'文件 {file_path} 准备处理...',
                'file_path': file_path,
                'file_type': file_ext,
                'knowledge_base': request.knowledge_base,
                'oss_url': oss_result.get('url', ''),
                'oss_success': oss_result['success']
            }
            
            # 异步启动处理任务
            import threading
            def async_process():
                process_file_async(file_path, file_ext, request.knowledge_base, task_id, oss_result.get('url'))
            
            thread = threading.Thread(target=async_process)
            thread.daemon = True
            thread.start()
        
        logger.info(f"路径上传成功，生成了 {len(task_ids)} 个任务")
        
        return UploadResponse(
            task_id=",".join(task_ids),  # 多个任务ID用逗号分隔
            message=f"成功准备处理 {len(request.file_paths)} 个文件，正在后台处理中。",
            status="uploaded",
            knowledge_base=request.knowledge_base
        )
        
    except Exception as e:
        logger.error(f"路径上传处理异常: {str(e)}")
        raise HTTPException(status_code=500, detail=str(e))

@router.get("/task/{task_id}", response_model=TaskStatusResponse, summary="查询任务状态")
async def get_task_status(task_id: str):
    """查询任务状态"""
    logger.info(f"查询任务状态: {task_id}")
    
    try:
        # 检查任务是否存在
        if task_id not in task_status:
            raise HTTPException(status_code=404, detail="任务不存在")
        
        task_info = task_status[task_id]
        
        return TaskStatusResponse(
            task_id=task_id,
            status=task_info['status'],
            progress=task_info.get('progress', 0),
            message=task_info.get('message', ''),
            file_path=task_info.get('file_path', ''),
            file_type=task_info.get('file_type', ''),
            knowledge_base=task_info.get('knowledge_base', ''),
            text_chunks_count=task_info.get('text_chunks_count', 0),
            images_count=task_info.get('images_count', 0),
            processing_time=task_info.get('processing_time', 0),
            error=task_info.get('error', ''),
            oss_url=task_info.get('oss_url', ''),
            oss_success=task_info.get('oss_success', False)
        )
            
    except Exception as e:
        logger.error(f"查询任务状态异常: {str(e)}")
        raise HTTPException(status_code=500, detail=str(e))

@router.get("/tasks", summary="获取所有任务状态")
async def get_all_tasks():
    """获取所有任务状态"""
    try:
        return {
            "total_tasks": len(task_status),
            "tasks": task_status
        }
    except Exception as e:
        logger.error(f"获取所有任务状态异常: {str(e)}")
        raise HTTPException(status_code=500, detail=str(e))

@router.get("/stats", response_model=UploadStats, summary="获取上传统计")
async def get_upload_stats():
    """获取文件上传统计信息"""
    try:
        # 统计任务状态
        total_tasks = len(task_status)
        completed_tasks = sum(1 for task in task_status.values() if task['status'] == settings.TASK_STATUS_COMPLETED)
        failed_tasks = sum(1 for task in task_status.values() if task['status'] == settings.TASK_STATUS_FAILED)
        processing_tasks = sum(1 for task in task_status.values() if task['status'] == settings.TASK_STATUS_PROCESSING)
        
        return UploadStats(
            total_files=total_tasks,
            processed_files=completed_tasks,
            failed_files=failed_tasks,
            processing_files=processing_tasks,
            total_size=0  # 可以添加文件大小统计
        )
    except Exception as e:
        raise HTTPException(status_code=500, detail=str(e))

@router.get("/supported-formats", summary="获取支持的文件格式")
async def get_supported_formats():
    """获取支持的文件格式列表"""
    return {
        "supported_formats": settings.ALLOWED_EXTENSIONS,
        "max_file_size": settings.MAX_FILE_SIZE,
        "max_file_size_mb": settings.MAX_FILE_SIZE // (1024 * 1024)
    } 